package com.nbsaas.boot.lock.impl;

import java.util.concurrent.*;
import java.util.concurrent.atomic.LongAdder;
import java.util.concurrent.locks.*;

public class ReentrantLockServiceWithCount extends ReentrantLockService {

    // 存储锁获取次数的计数器
    private final ConcurrentMap<String, LongAdder> lockCountMap = new ConcurrentHashMap<>();

    /**
     * 获取锁计数
     *
     * @param key 锁的唯一标识
     * @return 锁的获取次数
     */
    public int getLockCount(String key) {
        LongAdder count = lockCountMap.get(key);
        return count == null ? 0 : count.intValue();
    }

    @Override
    public boolean lock(String key) {
        boolean locked = super.lock(key);
        if (locked) {
            // 增加锁计数
            lockCountMap.computeIfAbsent(key, k -> new LongAdder()).increment();
        }
        return locked;
    }

    @Override
    public boolean lock(String key, long timeout, TimeUnit unit) {
        boolean locked = super.lock(key, timeout, unit);
        if (locked) {
            // 增加锁计数
            lockCountMap.computeIfAbsent(key, k -> new LongAdder()).increment();
        }
        return locked;
    }

    @Override
    public boolean unlock(String key) {
        boolean unlocked = super.unlock(key);
        if (unlocked) {
            // 解锁时也可以选择减少计数（如果需要）
            LongAdder count = lockCountMap.get(key);
            if (count != null) {
                count.decrement();
                if (count.intValue() == 0) {
                    lockCountMap.remove(key);
                }
            }
        }
        return unlocked;
    }

    /**
     * 获取所有锁的计数总数
     *
     * @return 所有锁的总计数
     */
    public int getTotalLockCount() {
        return lockCountMap.values().stream().mapToInt(LongAdder::intValue).sum();
    }
}
